#!/usr/bin/python

# agencyloader.py: Provides the AgencyLoader super-duper class

# Copyright 2011 Portland Transport

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# self.stops should be an array of GTFS Stop objects; they shall contain the
# code avl_id (even if it's the same as the stop_id or stop_tag). Remember,
# you (or your colleagues) have to access this data from Javascript.

# all _getStopsFrom... should return this format

# Authors:
# Matt Conway: main code

from urllib import urlretrieve
from couchdbkit import Database, Document, Server
import couchdbkit
from restkit import BasicAuth
from optparse import OptionParser
from urlparse import urlparse
import re

try:
    import json
except ImportError:
    print "JSON module not found, continuing without it. Reports will not work."

from time import sleep, gmtime, strftime
from zipfile import ZipFile
from sys import exc_info, exit
import csv
import os
from tempfile import mkdtemp
from shutil import rmtree
import subprocess

print 'Database load running as process %s' % os.getpid()
print strftime("Start: %a, %d %b %Y %H:%M:%S UTC", gmtime())

# Convenience code to replace sets
# Inspired by http://www.velocityreviews.com/forums/t355728-using-dictionaries-in-sets-dict-objects-are-unhashable.html,
# but heavily extended
def intersect(a, b):
    "Return the intersection of two lists"
    # Put in v in a if it's in b
    return [v for v in a if v in b]

def difference(a, b):
    "Returns the difference of two lists"
    return [v for v in a if not v in b]

class AgencyLoader:
    def __init__(self, options):
        self.options = options

    # We put it outside _getStopsFromGTFS so it can easily be overridden in a
    # subclass
    def _setAvlRouteID(self, route):
        route['avl_route_id'] = route['route_id']
        return route

    def _getStopsFromGTFS(self, gtfsurl, url=True):
        """Set url to false if you are passing in a file object. You can pass in
        either a URL or an open file object."""
        # The stop objects stop working if we don't keep a reference around
        global gtfs3925829405682_8375028573953

        tempdir = mkdtemp()

        if url:
            # download the file
            urlretrieve(gtfsurl, tempdir + '/gtfs.zip')
            gtfsurl = tempdir + '/gtfs.zip'

        # open the zip file
        gtfsZip = ZipFile(gtfsurl)
 
        # Make sure it's safe
        for theFile in gtfsZip.namelist():
            if '/' in theFile or '..' in theFile:
                print "Fatal: GTFS contains '/' or '..'"
                exit(1)
       
        # Can't use extract on python 2.5, and using read will cause the file 
        # to be loaded into memory, so use unzip
        subprocess.call(['unzip', gtfsurl, '-d', tempdir, 'stops.txt', 
                          'routes.txt', 'trips.txt', 'stop_times.txt',
                          'agency.txt', 'feed_info.txt'])

        stopsIn = csv.reader(open(tempdir + '/stops.txt'))
        routesIn = csv.reader(open(tempdir + '/routes.txt'))
        tripsIn = csv.reader(open(tempdir + '/trips.txt'))
        stopTimesIn = csv.reader(open(tempdir + '/stop_times.txt'))
        agenciesIn = csv.reader(open(tempdir + '/agency.txt'))

        # First line has column names
        stopsCols = stopsIn.next()
        routesCols = routesIn.next()
        tripsCols = tripsIn.next()
        stopTimesCols = stopTimesIn.next()
        agenciesCols = agenciesIn.next()

        self.agency = dict(zip(agenciesCols, agenciesIn.next()))
        if self.options.agency:
            self.agency['agency_name'] = self.options.agency

        del agenciesIn

        # Deal with agency_timezone being deprecated
        if not self.agency.has_key('agency_timezone'):
            feedInfoIn = csv.reader(open(tempdir + '/feed_info.txt'))
            feedInfoCols = feedInfoIn.next()

            if not 'feed_timezone' in feedInfoCols:
                raise KeyError, 'Fatal: timezone not found in either agency.txt or feed_info.txt!'

            feedInfoDict = dict(zip(feedInfoCols, feedInfoIn.next()))

            self.agency['agency_timezone'] = feedInfoDict['feed_timezone']

        print 'Loading GTFS for agency %s' % self.agency['agency_name']

        stops = {}
        for stop in stopsIn:
            stopDict = dict(zip(stopsCols, stop))
            stopDict['routes'] = []
            stops[stopDict['stop_id']] = stopDict
            
        del stopsIn

        print '%s stops in GTFS' % len (stops)
            
        # now run through the routes
        routes = {}
        for route in routesIn:
            routeDict = dict(zip(routesCols, route))
            try:
                del routeDict['agency_id']
            except KeyError:
                pass
            routes[routeDict['route_id']] = routeDict
            
        del routesIn

        print '%s routes in GTFS' % len (routes)
        
        # Build a mapping of trips to routes
        tripMap = {}
        for trip in tripsIn:
            tripDict = dict(zip(tripsCols, trip))
            tripMap[tripDict['trip_id']] = tripDict['route_id']
            
        del tripsIn

        print '%s trips in GTFS' % len (tripMap)
        
        for st in stopTimesIn:
            stDict = dict(zip(stopTimesCols, st))
            stop = stops[stDict['stop_id']]
            route = routes[tripMap[stDict['trip_id']]]
            if not route in stop['routes']:
                stop['routes'].append(route)

        del stopTimesIn

        rmtree(tempdir)

        print 'Done loading GTFS.'

        return stops.values()

    def loadAllStops(self):
        # Basic loader for GTFS-only agencies, sets avl_id same as
        # stop_id
        # Should be overwritten in subclasses
        stops = self._getStopsFromGTFS(self.options.gtfsUrl, not self.options.localgtfs)

        def setAvlIds(stop):
            stop['avl_stop_id'] = stop['stop_id']
            return stop
        self.stops = map(setAvlIds, stops)


    # Write an HTML report on what will be done
    def writeReport(self, outfile, addStops=[], delStops=[], newStopIds=[], 
                    delStopIds=[], updateStopIds=[]):
        # Build dictionaries
        addStops = dict([(stop['stop_id'], stop) for stop in addStops])
        delStops = dict([(stop['stop_id'], stop) for stop in delStops])

        # Make sure all the opt attrs are present
        def normalizeStop(stop):
            if not stop.has_key('stop_desc'):
                stop['stop_desc'] = ''

            if not stop.has_key('stop_url'):
                stop['stop_url'] = ''

            if not stop.has_key('stop_name'):
                stop['stop_name'] = ''

            if not stop.has_key('zone_id'):
                stop['zone_id'] = ''

            if not stop.has_key('avl_stop_id'):
                stop['avl_stop_id'] = ''
            
            return stop

        theJSON = {}
        # Only put in the ones that are actually new
        theJSON['new'] = [normalizeStop(addStops[stopId])
                          for stopId in newStopIds]
        theJSON['del'] = [normalizeStop(delStops[stopId])
                          for stopId in delStopIds]
        theJSON['updated'] = [dict(new=normalizeStop(addStops[stopId]),
                                   old=normalizeStop(delStops[stopId]))
                              for stopId in updateStopIds]
        
        # header
        outfile.write('''
<!DOCTYPE html>
<html>
<head>
  <title>Stop Loader Report</title>
  <script src="https://www.google.com/jsapi?key=ABQIAAAA2jbkDz9KQq1t-fKT2CS9MxQ9FV92uq6YW0A2bI7_oLN3w2-dWBRtH6W8sx6yVbZfTM-YOSc2CRGlLA" type="text/javascript"></script>
  <script src="https://ajax.googleapis.com/ajax/libs/jquery/1.6.2/jquery.min.js"
	  type="text/javascript"></script>
  <script type="text/javascript"
          src="http://maps.googleapis.com/maps/api/js?sensor=false"></script>
  <script src="http://transit-appliance-utils.googlecode.com/svn/trunk/NextBus/reports/report.js" type="text/javascript"></script>
  <link rel="stylesheet" href="report.css" type="text/css" />
  <script type="text/javascript">
     stops = ''')

        outfile.write(json.dumps(theJSON))

        outfile.write('''
</script>
</head>
<body>
  <!-- referencing the body itself to build the page is buggy in Chrome, so put
       everything in a container div -->
  <div id="mainContainer"></div>
</body>
</html>''')

        outfile.close()

    def pushToCouch(self):
        # fill out the stops
        # do it in place, for memory efficiency
        for stop in self.stops:
            stop['agency'] = self.agency['agency_name']
            stop['_id'] = '%s:%s' % (stop['agency'], stop['stop_id'])
            # Remove any None values so that the compare works.
            # CouchDB removes them automatically
            # This does not apply inside subattributes, for instance
            # values inside routes
            for key in stop.keys():
                if stop[key] == None:
                    del stop[key]

        # if there is a mapping function, run it
        if self.options.mapfunc:
            thePython = open(os.path.expanduser(self.options.mapfunc))
            code = compile(thePython.read(), self.options.mapfunc, 'exec')
            ns = {}
            exec code in ns
            
            # Now, run it.
            self.stops = [stop for stop in map(ns['stopMap'], self.stops) if stop != None]

        # Make sure it has an agency_name
        if self.options.agency:
            self.agency['agency_name'] = self.options.agency

        if not self.agency['agency_name']:
            print 'No agency_name set. Please override it with -a'
            exit(1)      

        # Check if auth needed        
        if self.options.dbuser is not None and self.options.dbpass is not None:
            print "Connecting to CouchDB as user %s" % self.options.dbuser
            # From http://support.cloudant.com/discussions/problems/59-couchdbkit-and-authentication
            server = Server(self.options.server, 
                            filters=[BasicAuth(self.options.dbuser,
                                      self.options.dbpass)])
        else:
            print 'Connecting to CouchDB anonymously'
            server = Server(self.options.server)

        # if the stopsDb is None, set it to the agency name (the default)
        # Database names must be lower case
        dbName = self.options.stopsDb or self.agency['agency_name'].lower()
        stopsDb = server.get_or_create_db(dbName)


        # We get the common elements, so we don't have to update them
        # The elements in the DB that aren't in the dataset should be deleted
        # The elements in the dataset that aren't in the DB should be added
        # This will also take care of updates b/c it will be both deleted and
        # added
        # thus, we delete first
        
        # Code to replace sets
        # found at http://www.velocityreviews.com/forums/t355728-using-dictionaries-in-sets-dict-objects-are-unhashable.html
        
        # Clear out all the Couch stuff (except _id) so that they compare
        # correctly. If anyone knows a better way to do this let me know
        # TODO: Improve.
        def removeCouchStuff(stop):
            for key in stop.keys():
                # We remove geometry because we don't want to compare on 
                # numerical latlon values which can be subject to rounding
                # errors - AC Transit has two such stops
                if (key[0] == '_' and key != '_id') or key == 'doc_type' or key == 'geometry':
                    del stop[key]
            return stop

        # Grab all of the stops for the agency, but don't grab the
        # agency defn or the _design documents
        # you can't just use stopsDb.documents due to a couchdbkit bug on
        # Dreamhost
        tempView = stopsDb.temp_view(
            {'map':'''function (doc) {
                         if (doc.agency && doc.agency == '%s') emit (doc._id, doc);
                      }''' % self.agency['agency_name']})

        dbStops = [removeCouchStuff(stop['value']) for stop in tempView.all()]

        print '%s stops in database' % len(dbStops)
        print '%s stops locally' % len(self.stops)

        common = intersect(self.stops, dbStops)

        toAdd = difference(self.stops, common)
        toDelete = difference(dbStops, common)

        if len(toAdd) == 0 and len(toDelete) == 0:
            print 'Agency %s is already loaded and up-to-date.' % self.agency['agency_name']
            print 'Nothing to do.'
            exit()

        # Calculate new stops, deleted stops and updated stops
        addIds = [stop['stop_id'] for stop in toAdd]
        delIds = [stop['stop_id'] for stop in toDelete]

        # Stops that are added but not deleted
        newStopIds = difference(addIds, delIds) 
        # Stops deleted but not added
        delStopIds = difference(delIds, addIds)
        # Stops both deleted and added
        updateStopIds = intersect(addIds, delIds)

        # Tell the user what we're going to do, then do it
        # Figure how many stops will be added, deleted, updated  

        while 1: # Loop in case the user gets command wrong, don't just exit
            print '%s stops to add, %s stops to delete, %s stops to update' % (len(newStopIds), len(delStopIds), len(updateStopIds))
            # We need to ask them if not
            if self.options.yes:
                break
            else:
                cmd = raw_input('Do you wish to continue (c), abort (a) or print a detailed report (r)? ')
                if cmd == 'a':
                    exit()
                elif cmd == 'c':
                    break
                elif cmd == 'r':
                    fn = raw_input('enter a file name (.html): ')
                    f = open(fn, 'w')
                    # No need to put in updateStops, b/c all updated stops are in
                    # toAdd and toDelete
                    self.writeReport(f, addStops=toAdd, delStops=toDelete, 
                                     newStopIds=newStopIds, delStopIds=delStopIds,
                                     updateStopIds=updateStopIds)
                    continue
                else:
                    continue

        # Add the agency first, so we never have data displayed sans-copyright
        # Just delete the old one.
        # We no longer use agenciesDb, as of 8/20/2011. Agencies are loaded 
        # directly into the same database as the stops.
        #agenciesDb = server.get_or_create_db(self.options.agenciesdb)

        self.agency['rights_notice'] = self.options.rights
        
        self.agency['_id'] = self.agency['agency_name']
    
        # Set it to the name
        self.agency['avl_agency_id'] = self.agency['agency_name']

        # Or the ID, if defined
        try: 
            self.agency['avl_agency_id'] = self.agency['avl_agency_id']
        except:
            pass

        # Or the override, if defined
        if self.options.avlagency:
            self.agency['avl_agency_id'] = self.options.avlagency

        # Allow it to be overridden
        try:
            self.agency['avl_service'] = self.options.avlservice or self.agency['avl_service']
        except KeyError:
            print 'Please set the AVL service.'
            exit(1)

        # Parse out the AVL options
        try:
            avlopts = self.agency['avl_options']
        except KeyError:
            # It wasn't set at all on the agency.
            avlopts = {}

        # Parse out all the AVL options - they're a list of strings
        # in the format key:value

        # Don't iterate if it's None
        if self.options.avloptions:
            for opt in self.options.avloptions:
                item = opt.split(':')
                if item[1] == 'true':
                    item[1] = True
                elif item[1] == 'false':
                    item[1] = False
                elif item[1] == 'null':
                    item[1] = None 
                    
                avlopts[item[0]] = item[1]

        self.agency['avl_options'] = avlopts

        try:
            stopsDb.delete_doc(self.agency['agency_name'])
        except couchdbkit.exceptions.ResourceNotFound:
            pass

        agencyDefn = Document(**self.agency)
        agencyDefn.set_db(stopsDb)
        agencyDefn.save()
    
        # Add the GeoCouch stuff now, to avoid rounding errors with the compare
        def addGeoCouch(stop):
            # GeoCouch
            stop['geometry'] = dict(
                # Yes, it goes longitude, latitude.
                # YOU ABSOLUTELY MUST USE A LIST!
                # b/c CouchDB returns only lists, and [] != (),
                # so for the comparison to work, 
                # YOU ABSOLUTELY MUST USE A LIST.
                coordinates=[float(stop['stop_lon']), float(stop['stop_lat'])],
                type='Point')
            
            stopdoc = Document(**stop)
            stopdoc.set_db(stopsDb)
            return stopdoc

        def primedb():
            # Seems like the database sometimes needs to be primed
            try:
                # Doesn't matter if it exists
                stopsDb.get('dshnfdkhjlkdsjgkldfjgsd')
            except:
                pass


        print 'Deleting %s stops' % len(toDelete)
        primedb()
        # Delete all in toDelete
        for stop in toDelete:
            delete = stopsDb.get(stop['_id'])
            for i in range(15):
                try:
                    stopsDb.delete_doc(delete)
                except:
                    print 'Failure, retrying %s' % i
                else:
                    break

        print 'Adding %s stops' % len(toAdd)
        primedb()
        for add in toAdd:
            for i in range(15):
                try:
                    stopsDb.save_doc(addGeoCouch(add))
                except:
                    'Failure, retrying %s' % i
                else:
                    break

        print '%s stops added/updated, %s stops deleted/updated!' % (len(toAdd), len(toDelete))

    def cleanup(self):
        print strftime("End: %a, %d %b %Y %H:%M:%S UTC", gmtime())

# You can pass in a parser that already has some options
# defined
def main(parser=OptionParser(), loaderClass=AgencyLoader):
    parser.add_option('-g', '--gtfs-url', dest='gtfsUrl', type='string',
                      help='The URL for the GTFS file.',
                      metavar='URL', default=None)

    parser.add_option('-a', '--agency', dest='agency', type='string',
                      help='Agency name.',
                      metavar='AGENCY')

    # Unfortunately, right now, couchdbkit/restkit only supports HTTP
    # See https://github.com/benoitc/couchdbkit/issues/issue/63
    # As of 2/11/11, the issue is closed, but I haven't tested with a new version of couchdbkit yet.
    parser.add_option('-s', '--server', dest='server', type='string',
                      help='The server that hosts the databases.',
                      metavar='URL',
                      default='http://127.0.0.1:5984/')

    parser.add_option('-d', '--database', dest='stopsDb', type='string',
                      help='The name of the stop database (default: agency name)',
                      metavar='DB', default=None)

    parser.add_option('-r', '--rights-notice', dest='rights',
                      type='string', 
                      help='Any notice that needs to be displayed with this agencies data', default='') # Set to blank string so it won't show as 'null' in HTML

    parser.add_option('-u', '--username', dest='dbuser', type='string',
                      help='The database username.',
                      metavar='USERNAME')

    parser.add_option('-p', '--password', dest='dbpass', type='string',
                      help='The database password.',
                      metavar='PASSWORD')

    parser.add_option('-y', '--yes', dest='yes', action='store_true',
                      help='Answer yes to prompts', default=False)

    parser.add_option('--local-gtfs', dest='localgtfs', action='store_true',
                      help='If you are using a local GTFS file',
                      default=False)

    parser.add_option('--avl-service', dest='avlservice',
                      help="The AVL service",
                      metavar='SERVICE', default=None)

    parser.add_option('--map-func', dest='mapfunc', type='string',
                      help='The file that contains a stop mapping function',
                      metavar='FILE', default=None)

    parser.add_option('--avl-option', dest='avloptions', action='append',
                      help='AVL option, in format KEY:VALUE.',
                      metavar='KEY:VAL')

    parser.add_option('--transform', dest='transform', type='string',
                      help='path or URL to a OneBusAway transform file',
                      metavar='TRANSFORM-FILE', default=None)

    # Allow them to specify the path to the transformer
    parser.add_option('--transformer-command', dest='transformer', 
                      help='The command for the transformer. Default: ' +
                      'java -jar gtfs-transformer.jar',
                      default='java -jar gtfs-transformer.jar')
                      

    parser.add_option('--avl-agency', dest='avlagency',
                      help='The AVL agency ID', metavar='AGENCY')

    options, args = parser.parse_args()

    al = loaderClass(options)
    al.loadAllStops()
    al.pushToCouch()
    al.cleanup()


if __name__ == '__main__':
    main()

    
